import os

from dotenv import load_dotenv
from langchain.chat_models import init_chat_model
from langchain.output_parsers import ResponseSchema, StructuredOutputParser
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate


def demo1():
    prompt = PromptTemplate.from_template(
        """请根据以下内容提取用户的信息，并以 JSON 格式返回：
    {input}
    返回格式如下：
    ```json
    {{  
      "name": "string",           // 用户的姓名
      "age": "string",            // 用户的年龄
      "year_of_birth": "string"   // 今年是2025年，根据年龄推算出生年
    }}
    ```"""
    )
    chain = (prompt | model | StrOutputParser())
    result = chain.invoke({"input": "我叫张三，今年18岁"})
    print(result)


def demo2():
    schemas = [
        ResponseSchema(name="name", description="用户的姓名"),
        ResponseSchema(name="age", description="用户的年龄"),
        ResponseSchema(name="year_of_birth", description="今年是2025年，根据用户的年龄推算他的出生年"),
    ]
    parser = StructuredOutputParser.from_response_schemas(schemas)
    prompt = PromptTemplate.from_template(
        "请根据以下内容提取用户的信息，并返回JSON格式:\n{input}\n\n{format_instructions}")
    # parser.get_format_instructions()
    # ```json
    # {
    #     "name": string // 用户的姓名
    #     "age": string // 用户的年龄
    #     "year_of_birth": string // 今年是2025年，根据用户的年龄推算他的出生年
    # }
    # ```
    chain = (prompt.partial(format_instructions=parser.get_format_instructions()) | model | parser)
    result = chain.invoke({"input": "我叫张三，今年18岁"})
    print(result)


if __name__ == '__main__':
    load_dotenv(override=True)
    DEEPSEEK_API_KEY = os.getenv("DEEPSEEK_API_KEY")
    model = init_chat_model(model="deepseek-chat", model_provider="deepseek")
    demo1()
